package ai.people.core.feign.ribbon;

import com.google.common.collect.Sets;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import lombok.extern.slf4j.Slf4j;

import java.net.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author yuanqinglong
 * @date 2022/3/31 12:45
 */
@Slf4j
public class LocalLoadBalancerRule implements IRule {

    private ILoadBalancer lb;


    @Override
    public Server choose(Object key) {
        // 1.获取到所有Server
        List<Server> servers = lb.getReachableServers();
        // 2.从所有Server中排除掉指定端口的Server后，剩余的Server
        List<Server> localServers = this.getLocalServers(servers);
        // 3.从剩余Server中随机选择一个Server
        int index = new Random().nextInt(localServers.size());
        return localServers.get(index);
    }

    /**
     * 获取本地服务器
     *
     * @param servers 服务器
     * @return {@link List}<{@link Server}>
     */
    private List<Server> getLocalServers(List<Server> servers) {
        try {
            Enumeration<NetworkInterface> allNetInterfaces = NetworkInterface.getNetworkInterfaces();
            Set<String> localIpSet = new HashSet<>();;
            while (allNetInterfaces.hasMoreElements()) {
                NetworkInterface netInterface = allNetInterfaces.nextElement();
                Enumeration<InetAddress> addresses = netInterface.getInetAddresses();
                while (addresses.hasMoreElements()) {
                    InetAddress inHostAddress = addresses.nextElement();
                    localIpSet.add(inHostAddress.getHostAddress());
                }
            }
            return servers.stream().filter(server -> localIpSet.contains(server.getHost())).collect(Collectors.toList());
        } catch (Exception e) {
            log.warn("自定义负载均衡策略获取本地IP异常", e);
            return servers;
        }
    }




    @Override
    public void setLoadBalancer(ILoadBalancer lb) {
        this.lb = lb;
    }

    @Override
    public ILoadBalancer getLoadBalancer() {
        return lb;
    }
}
